{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install -qU memorizz yahooquery"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import getpass\n",
    "import os\n",
    "\n",
    "# Function to securely get and set environment variables\n",
    "def set_env_securely(var_name, prompt):\n",
    "    value = getpass.getpass(prompt)\n",
    "    os.environ[var_name] = value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "set_env_securely(\"MONGODB_URI\", \"Enter your MongoDB URI: \")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "set_env_securely(\"OPENAI_API_KEY\", \"Enter your OpenAI API Key: \")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from memorizz.memory_provider.mongodb.provider import MongoDBConfig, MongoDBProvider\n",
    "\n",
    "# Create a memory provider\n",
    "mongodb_config = MongoDBConfig(uri=os.environ[\"MONGODB_URI\"])\n",
    "memory_provider = MongoDBProvider(mongodb_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "from memorizz.memory_component.memory_mode import MemoryMode\n",
    "from memorizz.workflow.workflow import Workflow\n",
    "from memorizz import Toolbox\n",
    "import requests\n",
    "\n",
    "def get_weather(latitude, longitude):\n",
    "    response = requests.get(f\"https://api.open-meteo.com/v1/forecast?latitude={latitude}&longitude={longitude}&current=temperature_2m,wind_speed_10m&hourly=temperature_2m,relative_humidity_2m,wind_speed_10m\")\n",
    "    data = response.json()\n",
    "    return data['current']['temperature_2m']\n",
    "\n",
    "def calculate_distance(city1: str, city2: str) -> str:\n",
    "    \"\"\"Calculate distance between two cities.\"\"\"\n",
    "    # Simulate distance calculation\n",
    "    distances = {\n",
    "        (\"New York\", \"London\"): \"3,459 miles (5,567 km)\",\n",
    "        (\"New York\", \"Tokyo\"): \"6,743 miles (10,850 km)\",\n",
    "        (\"London\", \"Paris\"): \"214 miles (344 km)\",\n",
    "        (\"Tokyo\", \"Sydney\"): \"4,840 miles (7,788 km)\",\n",
    "        (\"New York\", \"Paris\"): \"3,631 miles (5,844 km)\"\n",
    "    }\n",
    "    key = (city1, city2) if (city1, city2) in distances else (city2, city1)\n",
    "    return distances.get(key, f\"Distance data not available between {city1} and {city2}\")\n",
    "\n",
    "def book_flight(origin: str, destination: str, date: str) -> str:\n",
    "    \"\"\"Book a flight between two cities.\"\"\"\n",
    "    # Simulate flight booking (will sometimes \"fail\" for demonstration)\n",
    "    if \"Mars\" in origin or \"Mars\" in destination:\n",
    "        raise Exception(\"No flights available to Mars at this time\")\n",
    "    \n",
    "    price = hash(f\"{origin}{destination}{date}\") % 500 + 200\n",
    "    return f\"Flight booked from {origin} to {destination} on {date}. Price: ${price}. Confirmation: FL{hash(date) % 10000}\"\n",
    "\n",
    "def get_currency_rate(from_currency: str, to_currency: str) -> str:\n",
    "    \"\"\"Get currency exchange rate.\"\"\"\n",
    "    # Simulate currency rates\n",
    "    rates = {\n",
    "        (\"USD\", \"EUR\"): \"0.85\",\n",
    "        (\"USD\", \"GBP\"): \"0.73\",\n",
    "        (\"USD\", \"JPY\"): \"110.0\",\n",
    "        (\"EUR\", \"USD\"): \"1.18\",\n",
    "        (\"GBP\", \"USD\"): \"1.37\"\n",
    "    }\n",
    "    key = (from_currency.upper(), to_currency.upper())\n",
    "    rate = rates.get(key, \"1.00\")\n",
    "    return f\"1 {from_currency.upper()} = {rate} {to_currency.upper()}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create toolbox and register tools, with memory provider\n",
    "toolbox = Toolbox(memory_provider=memory_provider)\n",
    "\n",
    "toolbox.register_tool(get_weather)\n",
    "toolbox.register_tool(calculate_distance)\n",
    "toolbox.register_tool(book_flight)\n",
    "toolbox.register_tool(get_currency_rate)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def print_section(title):\n",
    "    \"\"\"Print a formatted section header.\"\"\"\n",
    "    print(f\"\\n{'='*60}\")\n",
    "    print(f\" {title}\")\n",
    "    print(f\"{'='*60}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def print_workflow_info(workflow):\n",
    "    \"\"\"Print workflow information in a readable format.\"\"\"\n",
    "    print(f\"Workflow: {workflow.name}\")\n",
    "    print(f\"Description: {workflow.description}\")\n",
    "    print(f\"Outcome: {workflow.outcome.value}\")\n",
    "    print(f\"Created: {workflow.created_at}\")\n",
    "    print(f\"Steps: {len(workflow.steps)}\")\n",
    "    \n",
    "    for step_name, step_data in workflow.steps.items():\n",
    "        print(f\"  - {step_name}\")\n",
    "        print(f\"    Function: {step_data.get('_id', 'Unknown')}\")\n",
    "        print(f\"    Arguments: {step_data.get('arguments', {})}\")\n",
    "        if step_data.get('error'):\n",
    "            print(f\"    Error: {step_data.get('error')}\")\n",
    "        else:\n",
    "            result = step_data.get('result', 'No result')\n",
    "            # Handle different result types properly\n",
    "            if isinstance(result, str):\n",
    "                print(f\"    Result: {result[:100]}...\" if len(result) > 100 else f\"    Result: {result}\")\n",
    "            else:\n",
    "                # For non-string results (like float, int, etc.), convert to string first\n",
    "                result_str = str(result)\n",
    "                print(f\"    Result: {result_str[:100]}...\" if len(result_str) > 100 else f\"    Result: {result_str}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "from memorizz.memagent import MemAgent\n",
    "\n",
    "def demonstrate_workflow_memory():\n",
    "    \"\"\"Main demonstration function.\"\"\"\n",
    "    \n",
    "    print_section(\"Workflow Memory Demonstration\")\n",
    "    print(\"This demo shows how MemAgent learns from previous tool executions\")\n",
    "    print(\"and uses that knowledge to guide future similar tasks.\\n\")\n",
    "    \n",
    "    # Create memory provider and toolbox    \n",
    "    # Create agent with workflow memory mode\n",
    "    print(\"Creating MemAgent with Workflow memory mode...\")\n",
    "    agent = MemAgent(\n",
    "        tools=toolbox,\n",
    "        memory_mode=MemoryMode.Workflow, # Set the memory mode to Workflow\n",
    "        memory_provider=memory_provider, # Add the memory provider\n",
    "        instruction=\"You are a helpful travel assistant. Use the available tools to help users with travel-related queries.\"\n",
    "    )\n",
    "\n",
    "    # agent.add_tool(toolbox=toolbox)\n",
    "\n",
    "    print(\"Created agent with tools: \", agent.tools)\n",
    "    \n",
    "    # Save the agent\n",
    "    agent.save()\n",
    "    \n",
    "    # Phase 1: Create initial workflows\n",
    "    print_section(\"Phase 1: Creating Initial Workflows\")\n",
    "    \n",
    "    # Successful workflow example\n",
    "    print(\"\\n1. Running travel planning query (will create successful workflow)...\")\n",
    "    response1 = agent.run(\"I want to plan a trip from New York to London. Can you help me with weather and distance information?\")\n",
    "    print(f\"Response: {response1}...\")\n",
    "    \n",
    "    time.sleep(1)  # Brief pause for demo effect\n",
    "    \n",
    "    # Failed workflow example\n",
    "    print(\"\\n2. Running impossible travel query (will create failed workflow)...\")\n",
    "    try:\n",
    "        response2 = agent.run(\"I want to book a flight from New York to Mars for next week, what is the weather like in Mars?\")\n",
    "        print(f\"Response: {response2}...\")\n",
    "    except Exception as e:\n",
    "        print(f\"Query completed with some failures as expected: {str(e)}...\")\n",
    "    \n",
    "    time.sleep(1)\n",
    "    \n",
    "    # Another successful workflow\n",
    "    print(\"\\n3. Running currency conversion query (will create successful workflow)...\")\n",
    "    response3 = agent.run(\"What's the exchange rate from USD to EUR, and what's the weather like in Paris?\")\n",
    "    print(f\"Response: {response3[:200]}...\")\n",
    "    \n",
    "    # Phase 2: Show workflow retrieval\n",
    "    print_section(\"Phase 2: Examining Created Workflows\")\n",
    "    \n",
    "    # Retrieve workflows to show what was stored\n",
    "    print(\"Retrieving stored workflows...\")\n",
    "    workflows = Workflow.retrieve_workflows_by_query(\"travel planning\", memory_provider)\n",
    "    \n",
    "    print(f\"\\nFound {len(workflows)} workflows related to travel planning:\")\n",
    "    for i, workflow in enumerate(workflows, 1):\n",
    "        print(f\"\\n--- Workflow {i} ---\")\n",
    "        print_workflow_info(workflow)\n",
    "    \n",
    "    # Phase 3: Demonstrate workflow memory in action\n",
    "    print_section(\"Phase 3: Workflow Memory Guidance in Action\")\n",
    "    \n",
    "    print(\"Now running a similar query to see how workflow memory guides execution...\")\n",
    "    print(\"The agent should reference previous successful workflows for travel planning.\\n\")\n",
    "    \n",
    "    # This query should benefit from previous workflow memory\n",
    "    response4 = agent.run(\"I'm planning another trip from London to Tokyo. Can you help me with the same kind of information as before?\")\n",
    "    print(f\"Response: {response4[:300]}...\")\n",
    "    \n",
    "    print(\"\\n\" + \"=\"*60)\n",
    "    print(\"Notice how the agent can reference previous workflow executions\")\n",
    "    print(\"to understand what 'same kind of information' means and follow\")\n",
    "    print(\"similar successful patterns from past workflows.\")\n",
    "    print(\"=\"*60)\n",
    "    \n",
    "    # Phase 4: Show workflow failure learning\n",
    "    print_section(\"Phase 4: Learning from Failures\")\n",
    "    \n",
    "    print(\"Running another impossible travel query to show failure learning...\")\n",
    "    try:\n",
    "        response5 = agent.run(\"Can you book me a flight to Mars? I know it didn't work before.\")\n",
    "        print(f\"Response: {response5[:200]}...\")\n",
    "    except Exception as e:\n",
    "        print(f\"Query completed: {str(e)[:100]}...\")\n",
    "    \n",
    "    print(\"\\nThe agent should learn from previous failures and potentially\")\n",
    "    print(\"avoid repeating the same failed approaches.\")\n",
    "    \n",
    "    # Final summary\n",
    "    print_section(\"Workflow Memory Benefits Demonstrated\")\n",
    "    print(\"✓ Tool executions are tracked as workflows\")\n",
    "    print(\"✓ Successful workflows guide future similar tasks\")\n",
    "    print(\"✓ Failed workflows help avoid repeating mistakes\")\n",
    "    print(\"✓ Context from past executions improves task understanding\")\n",
    "    print(\"✓ Learning accumulates over time for better performance\")\n",
    "    \n",
    "    print(f\"\\nAgent ID for future reference: {agent.agent_id}\")\n",
    "    print(\"You can load this agent later to see accumulated workflow memory.\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    demonstrate_workflow_memory()\n",
    "except KeyboardInterrupt:\n",
    "    print(\"\\nDemo interrupted by user.\")\n",
    "except Exception as e:\n",
    "    print(f\"\\nDemo error: {e}\")\n",
    "    import traceback\n",
    "    traceback.print_exc()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "memorizz",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
